import numpy as np
import random
import value_net

trainNum = 0
winList = []
loseList = []

# warn: probMap和posList用makeCompleteProbMap合成，（合成后的）probMap和board用toModelBoard合成

class situation:
    def __init__(self, board:list, posList:list, otherFeature:list, isWin:bool):
        self.board = np.array(board)
        self.posList = posList
        self.otherFeature = np.array(otherFeature)

        if isWin:
            winList.append(self)
        else:
            loseList.append(self)

def train(modelObj, epoch:int, batch_size:int, totEpoch:int):
    for _ in range(totEpoch):
        random.shuffle(winList)
        random.shuffle(loseList)
        total_size=min(len(winList),len(loseList))

        allBoard = []
        allOtherFeature = []
        allIsWin = []

        def add(sample, isWin):
            newBoard = value_net.toModelBoard(sample.board, sample.probMap)
            allBoard.append(newBoard)
            allOtherFeature.append(sample.otherFeature)
            allIsWin.append(isWin)

        for i in range(total_size): # 每次训练里面抽total_size个(winNum+loseNum)/batch_size个)，全过一遍
            add(winList[i], 1)
            add(loseList[i], 0)

        allBoard = np.array(allBoard)
        allOtherFeature = np.array(allOtherFeature)
        allIsWin = np.array(allIsWin)
        modelObj.train(allBoard, allOtherFeature, allIsWin, epoch, batch_size)

        global trainNum
        modelObj.save_model('model'+str(trainNum)+'.pkl')
        trainNum += 1
